%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Of Cities and Slums: Main Code for Solving Counterfactuals in 2010
% % ---- Calls the Functions: 
%     RuralEquilibriumFunction.m,
%     M_OnlyFunction.m and 
%     UrbanEquilibriumFunction.m
%     EquilibriumFunction.m
%     EquilibriumFunction_RandomAssign_Counterfactual.m
%     EquilibriumFunction_SelectiveAssign_Counterfactual.m

%% Outline
% I.      Setting Up the Economy: Parameters and Functional Forms
% II.     Setting Up the Code: Computational Aspects, e.g. gridpoints, and options.
% III.    Initial Conditions: Values Set for population locations, etc
% IV.     Big Loop: Equilibrium location decisions, q_R, q_F, q_C
%           Inner Loop: Equilibrium Prices and Quantities, given locations.
% V.      Reporting the Results      


clear all
close all
warning off

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%      Initial Conditions
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%       
%% Population
    %Pop_data = [0.2874; 0.10]; %1980 25-40
    Pop_data = [0.1368; 0.1899]; %2010 25-40

%% Education (Schooling) Levels (average education in each bin: 0, 1-4, 5-8, 9-11, 12+)
    %EE_set= [0, 3.17, 6.90, 10.72, 14.79]';  % Using Brazilian data (1991)
    EE_set= [0, 3.16, 6.98, 10.71, 14.68]';  % Using Brazilian data (1980)

%% Initial State: cross-section education for the country
    %pi_E = [0.2512 0.4550 0.1294 0.0955 0.0689]; %1980 25-40
    pi_E = [0.0916 0.3074 0.2742 0.2397 0.0871]; %2000 25-40


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%      Computational Aspects
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    %% Grid Sizes for Labor Market Skills
        Nxr = 200/2;
        Nxn = 300/2;

    %% Grid for numerical integration
        uB_xr        =    8; % 2*max(z_b*sigma_r_e + mu_r_e);
        uB_xn        =    8; % 2*max(z_b*sigma_n_e + mu_n_e);
        lB_xr        =   -10; %2*min(-z_b*sigma_r_e + mu_r_e);
        lB_xn        =   -10; %2*min(-z_b*sigma_r_e + mu_r_e);

        xr=linspace(lB_xr,uB_xr,Nxr);
        xn=linspace(lB_xn,uB_xn,Nxn)';
        ones_xr=ones(size(xr));
        ones_xn=ones(size(xn));

        xrxr=kron(xr,ones_xn);
        xnxn=kron(xn,ones_xr);

        Xr=exp(xr);
        XrXr=exp(xrxr);
        Xn=exp(xn);
        XnXn=exp(xnxn);

 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%      Parameters and Functional Forms
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%   

%% Labor Market Skills:
% --Previously generated to match cross-section distribution of earnings across occupations and education groups in Brazil
      NE=5; %number of educational groups               
      Moments_Skills_e = [   -1.2224   -7.3039    0.6669    2.7420   -0.0180
                            0.2326   -4.6167    0.7264    2.4773    0.1971
                            0.5676   -2.6523    0.7762    1.9710    0.3421
                            0.6784   -2.2963    0.9511    2.2052   -0.6787
                            1.2115    1.1570    0.8432    1.1728    0.4740];                  
        mu_r_e    = Moments_Skills_e(:,1);
        mu_n_e    = Moments_Skills_e(:,2);
        sigma_r_e = Moments_Skills_e(:,3);
        sigma_n_e = Moments_Skills_e(:,4);
        rho_rn_e  = Moments_Skills_e(:,5);

% A Bivariate Log-normal Skills Distribution for each e \in E 
    for e=1:NE,    
        mu_xr    =   mu_r_e(e);
        mu_xn    =   mu_n_e(e);
        sigma_xr =   sigma_r_e(e);
        sigma_xn =   sigma_n_e(e);
        Sigma_xr =   sigma_xr^2;
        Sigma_xn =   sigma_xn^2;
        rho_rn   =   rho_rn_e(e);
        cov_xrxn =   rho_rn*sigma_xr*sigma_xn;

        frn_e(:,:,e)       = (1/(2*pi*(1-rho_rn^2)^.5*sigma_xr*sigma_xn))*exp( -1/(2*(1-rho_rn^2))*( (xrxr-mu_xr).^2/Sigma_xr + (xnxn-mu_xn).^2/Sigma_xn   -2*rho_rn/sigma_xr/sigma_xn*(xrxr-mu_xr).*(xnxn-mu_xn) ));
        fXrXn_e(:,:,e)     = frn_e(:,:,e)./XrXr./XnXn;  % densities for the exponentiated variables
        fXr_e(:,:,e)   = frn_e(:,:,e)./XnXn;
        fXn_e(:,:,e)   = frn_e(:,:,e)./XrXr;
end


%% Estimated Parameters for the Intergenerational Mobility in Skill Formation:
% --Call the function TransProb.m
% Constructing the Implied Transition Functions
    paramsC =  [    0.3752    0.0602   -0.2782    0.0603   -1.4588    0.1549    2.2017   -0.8731
                    0.3452    0.2609   -0.5171    0.4215    0.6049    0.2326   -0.6450    0.3876
                    0.5579    0.0578    2.6716   -0.2078    3.5317   -0.1805    1.8001    0.0899
                    -0.0527    0.2080   -0.3673    0.3916   -0.7364    0.7135   -0.9283    0.7899
                    -2.9930    1.3062   -1.5812    1.2301   -1.2127    1.3262   -1.7434    1.5541];            

    paramsF =  [   -1.3926    0.3519   -4.3194    0.8462    1.7253   -2.2400    5.6594   -6.0559
                    0.3826    0.1879   -1.3421    0.4963   -0.8398    0.1813    0.7335   -0.9144
                    0.1229    0.1407   -0.2644    0.3590    0.1099    0.1168   -0.3958   -0.2198
                    0.3636    0.0828    0.8828    0.0198   -0.1174    0.4303   -0.3179    0.1135
                    -0.5405    1.0020   -0.7428    1.0815   -0.8126    1.2632   -0.5466    1.0847];
           
    paramsU =  [   -1.3447    0.3221   -4.0053    0.6585   -6.5037    0.9764   -0.2165   -0.5487
                    1.5671   -0.0174    0.7875    0.1079    1.9935   -0.1216    1.2530   -0.0760
                    2.1359   -0.0939    1.4595    0.0846    3.2487   -0.0719    1.8383    0.1454
                    1.0022    0.0308    2.5224   -0.0673    3.3225   -0.0551    2.8724    0.0567
                    2.9370   -0.0408    4.1377   -0.0805    4.3052    0.0418    3.8181    0.2640];
            

    paramsR =   [   -1.4031    0.5588   -3.7842    0.7317 -153.0621    0.4551 -343.9319    0.6787
                    0.2595    0.2825   -2.5330    1.0551    3.1706   -3.5829  -11.6147    2.9642
                    -5.0708    1.8372   -2.3179    0.9077    0.1668   -0.1033   -1.9433    0.5911
                    2.4999   -0.0915    1.4746    0.1801    3.4392   -0.4036    1.8471    0.2782
                    0.0341   -0.0626    2.8190   -0.9117    3.5031   -1.0195   -0.9551    0.6864];


%% Housing Prices: location l,   
% Housing Cost Options:
    %--- HousingCostOption=1 :  xiH_l    =   xiH0_l*(1+pop_l)^xiH1_l;
    %--- HousingCostOption=2 :  xiH_l    =   xiH0_l*(pop_l)^xiH1_l;
    %--- HousingCostOption=3 :  xiH_l    =   xiH0_l* exp( xiH1_l*(pop_l));
    HousingCostOption=1;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%Initial guess
alphaS=0.159580914;
cSbar=0.204026476;
cAbar=1.627291831;
beta=0.208430069;
gamE=0.002875338;
etaE=2.345172937;  
tau_F=0.00553139810140966;
xiH0_F=0.337370544;
xiH0_C=1.333518769;
XA=25.4928983029559;
XS=5.52456685463206;
XM=7.88081875022141;
xiH1_F=0.23;
xiH1_C=0.3;
varphi=0.54;
alphaA=0.01;
tau_F_income=0.16;
gam=1;
tau_R=0;
tau_C=0;
tau_1_F=0;

parameter_vector_all = [alphaS, cSbar, cAbar, beta, gamE, etaE, tau_F, xiH0_F, xiH0_C, XA, XS, XM, xiH1_F, xiH1_C, varphi, alphaA, tau_F_income,gam, tau_R, tau_C, tau_1_F];

%% Parameter to choose the counterfactual exercise:
    counterfactual_choice = 8;
    % 1: Benchmark economy
    % 2: City-schools in rural areas
    % 3: City-schools in slums
    % 4: Random assignment, i.e., for a probability say theta, 1-theta of the slum kids stay in the slum and theta are added to the school.
    % 5: Selective assignment (9+), i.e. the probabilities are e-dependent, that is  the probability of being sent to the C-schools are increasing in the education of the parents.
    % 6: Selective assignment (1+), i.e. the probabilities are e-dependent, that is  the probability of being sent to the C-schools are increasing in the education of the parents.
    % 7: Eliminate slums (tau_f_income=0.99)
    % 8: Facilitate slums (tau_f_income=0.0)

    if counterfactual_choice == 1
        % Benchmark economy
        equilibrium = EquilibriumFunction(parameter_vector_all,Pop_data,pi_E,EE_set,HousingCostOption,NE,Nxr,Nxn,paramsR,paramsF,paramsC,xr,xn,Xn,Xr,XrXr,XnXn,frn_e);
    
    elseif counterfactual_choice == 2
        % City-schools in rural areas
        paramsR = paramsC;
        equilibrium = EquilibriumFunction(parameter_vector_all,Pop_data,pi_E,EE_set,HousingCostOption,NE,Nxr,Nxn,paramsR,paramsF,paramsC,xr,xn,Xn,Xr,XrXr,XnXn,frn_e);
    
    elseif counterfactual_choice == 3
        % City-schools in slums
        paramsF = paramsC;
        equilibrium = EquilibriumFunction(parameter_vector_all,Pop_data,pi_E,EE_set,HousingCostOption,NE,Nxr,Nxn,paramsR,paramsF,paramsC,xr,xn,Xn,Xr,XrXr,XnXn,frn_e);
    
    elseif counterfactual_choice == 4
        % Random assignment, i.e., for a probability say theta, 1-theta of the slum kids stay in the slum and theta are added to the school.
        index_educ=0.75;
        equilibrium = EquilibriumFunction_RandomAssign_Counterfactual(parameter_vector_all,Pop_data,pi_E,EE_set,HousingCostOption,NE,Nxr,Nxn,paramsR,paramsF,paramsC,xr,xn,Xn,Xr,XrXr,XnXn,frn_e,index_educ);
    
    elseif counterfactual_choice == 5
        %Selective assignment, i.e. the probabilities are e-dependent, that is  the probability of being sent to the C-schools are increasing in the education of the parents.
        index_educ=2;
        equilibrium = EquilibriumFunction_SelectiveAssign_Counterfactual(parameter_vector_all,Pop_data,pi_E,EE_set,HousingCostOption,NE,Nxr,Nxn,paramsR,paramsF,paramsC,xr,xn,Xn,Xr,XrXr,XnXn,frn_e,index_educ);

    elseif counterfactual_choice == 6
        %Selective assignment, i.e. the probabilities are e-dependent, that is  the probability of being sent to the C-schools are increasing in the education of the parents.
        index_educ=4;
        equilibrium = EquilibriumFunction_SelectiveAssign_Counterfactual(parameter_vector_all,Pop_data,pi_E,EE_set,HousingCostOption,NE,Nxr,Nxn,paramsR,paramsF,paramsC,xr,xn,Xn,Xr,XrXr,XnXn,frn_e,index_educ);
    
    elseif counterfactual_choice == 7
        %Eliminate slums - labor Costs: tau_F_income=0.99
        tau_F_income=0.99;
        parameter_vector_all = [alphaS, cSbar, cAbar, beta, gamE, etaE, tau_F, xiH0_F, xiH0_C, XA, XS, XM, xiH1_F, xiH1_C, varphi, alphaA, tau_F_income,gam, tau_R, tau_C, tau_1_F];
        equilibrium = EquilibriumFunction(parameter_vector_all,Pop_data,pi_E,EE_set,HousingCostOption,NE,Nxr,Nxn,paramsR,paramsF,paramsC,xr,xn,Xn,Xr,XrXr,XnXn,frn_e);

    elseif counterfactual_choice == 8
        %Facilitate slums - labor Costs: tau_F_income=0.0
        tau_F_income=0.0;
        parameter_vector_all = [alphaS, cSbar, cAbar, beta, gamE, etaE, tau_F, xiH0_F, xiH0_C, XA, XS, XM, xiH1_F, xiH1_C, varphi, alphaA, tau_F_income,gam, tau_R, tau_C, tau_1_F];
        equilibrium = EquilibriumFunction(parameter_vector_all,Pop_data,pi_E,EE_set,HousingCostOption,NE,Nxr,Nxn,paramsR,paramsF,paramsC,xr,xn,Xn,Xr,XrXr,XnXn,frn_e);
    end

     

